
% This code is written by Dr. Hungtang Ko.
% For a manuscript titled "Collective movement of schooling fish reduces locomotor cost in turbulence".
% Last modified on 4/17/2024
% This code extracts the midline time series from the video. 

%% initialization 
close all; clear; clc
global flag_fig size_bound

% Flag for generating figure 
flag_fig = false; 

% Size limit for excluding abnormal regions 
% size_bound = [1500,3000];% for schools
size_bound = [300,8000];% % for individuals

%% read video

% names of the folders to be processed
folders_laminar = "Danio-schooling laminar/"+[
    {'Danio.School2N=8.Feb21.2022'                     }
    {'Danio.School3N=8.Feb17.2022'                     }
    {'Danio.School4N=8.Feb19.2022'                     }
    {'Danio.School5N=8.Feb23.2022'                     }];
folders_turbulent = "Danio-schooling turbulence/"+[
    {'Turbulence video-School4-June 29 2022'    }
    {'Turbulence video_School1-July 8 2022'     }
    {'Turbulence video_School2_July 19 2022'    }
    {'Turbulence video_School5-July 14 2022'    }];
folders_laminar_individual = "Danio-schooling laminar/"+[
    {'Daino.Individual1-March16.2022'                     }
    {'Danio.Individual2-March18.2022'                     }
    {'Danio.Individual3-March20.2022'                     }
    {'Danio.Individual4-March22.2022'                     }
    {'Danio.Individual5.March24.2022'                     }];
folders_turbulent_individual = "Danio-schooling turbulence/"+[
    {'Turbulence video_Individual1_July 17 2022'    }
    {'Turbulence Video_Individual2_July 6 2022'     }
    {'Turbulence video_Individual3_July 12 2022'    }];
data = [];

% Folder = folders_laminar;
% Folder = folders_turbulent;
% Folder = folders_laminar_individual;
Folder = folders_turbulent_individual;

%% for each folder & file
for i_folder = 1:length(Folder)
folder = Folder(i_folder);

% in the i_trial th folder
filenames = dir(folder);
filenames = {filenames.name};
filenames = filenames(contains(filenames,'avi'));
filenames = filenames(contains(filenames,'C2'))';
filenames = filenames(~contains(filenames,'._'))';

for i_case = 1:length(filenames)

    % display the file name being processed
    fn_long = folder+"/"+filenames(i_case);
    disp("Analyzing "+fn_long)
    
    % read video
    vid = VideoReader(fn_long);
    
    % generate background
    Background = background_gen(vid,1000,0);
    
    %% Main processing loop 
    % for a uniformly spaced starting timestamps
    for tstart = (0:0.5:9.5)*125/30
    
        %% auto generate suggestions for each vid and frame
        
        % set the video current time to the timestamp
        try 
            vid.CurrentTime = tstart;
        catch
            disp("video too short")
            break
        end

        % read and proces frames
        vidFrame = readFrame(vid);
        Diff = double(vidFrame)/255-Background;
        % BW = Diff>0.05; % for schools
        BW = vidFrame>0.5*255; % for individuals
        BW = bwareaopen(BW,1000); % remove area too small
        % s  = regionprops(BW, 'area', 'centroid');
        CC = bwconncomp(BW);
        CC = CC.PixelIdxList;
        tip = [];

        % for all patches that satisfy the the size condition, find the tips of the fish 
        for i = 1:length(CC)
        if length(CC{i}) > size_bound(1) && length(CC{i}) < size_bound(2)  
            y_temp = (mod(CC{i}(1),vid.height));
            x_temp = floor(CC{i}(1)/vid.height)+30;
            tip = [tip;x_temp,y_temp];
        end
        end
        
        % generate a figure
        if flag_fig
            figure(1)
            subplot(2,1,1)
            imshowpair(vidFrame,BW);
            subplot(2,1,2)
            if ~isempty(tip)
                imshowpair(vidFrame,bwselect(BW,tip(:,1),tip(:,2)));
            end
        end
        
        % process an instance/fish starting from the tip location
        for i_instance = 1:size(tip,1)
            stat = midline_individual(vid,Background,tstart,tip(i_instance,1),tip(i_instance,2));
            
            % plotting the instance
            if flag_fig && ~isempty(stat)
                figure(2)
                a = subplot(2,2,1);
                a.Position(3) = a.Position(3)*2;
                a.Position(1) = (1-(a.Position(3)))/2; 
                imshow(stat.BW)
                subplot(2,2,3)
                plot((1:250)/stat.length,stat.midline)
                axis equal
                subplot(2,2,4)
                hold all
                plot(stat.midline(:,1),'b')
                plot(stat.midline(:,stat.length-10),'r')
                legend('nose','tail','location','southeast')
            end
            
            % saving, simple exclusion
            if ~isempty(stat) && size(stat.midline,1)>10 % at least have 10 data points ~ 1 cycle
                data = [data,stat];
                disp('instance saved')
            end    
        end
    end
end
end

%% saving
save("data/midline/midlines_turbulent_individual",'data')



%% functions

% This  function tracks a fish starting from a specified nose location at a key frame
% Its output variable is a structure named "stat", which include the following fields
%   vid: video name
%   path: video path
%   tstart: the timestamp of the keyframe 
%   width: width of the fish
%   length: length of the fish
%   midline: midline of the fish in body lengths
%   amplitude_nose: nose amplitude
%   amplitude_tail: tail amplitude 

function stat = midline_individual(vid,Background,tstart,x_input,y_input)
global flag_fig size_bound

stat.vid = vid.Name;
stat.path = vid.Path;
stat.tstart = tstart;

vid.CurrentTime = tstart; 

Midline = [];
Width = [];
while 1
    try 
        vidFrame2 = readFrame(vid);
        Diff = double(vidFrame2)/255-Background;
        BW = imbinarize(Diff,0.05); % independent of school/individual condition
    catch 
        BW = false(vid.Width,vid.Height);
    end
    BW1 = bwselect(BW,x_input,y_input); 

%     sum(BW1,"all")
%     pause(0.1)
    if sum(BW1,"all") > size_bound(2)|| sum(BW1,"all") < size_bound(1)
        if flag_fig
            figure(1)
            sum(BW1,"all")
            imshowpair(BW,BW1)
        end
        if flag_fig
            pause
        end
        break
    end

    [y,x] = find(BW1);
    x_nose = min(x);
    y_nose = min(y(x == x_nose));

    x_midline = unique(x);
    if length(x_midline)<250    
        x_midline = padarray(x_midline,250-length(x_midline),nan,'post');
    else
        x_midline = x_midline(1:250);
    end
    y_midline = nan(1,250);
    width = nan(1,250);    

    for j = 1:length(x_midline)
    if ~isnan(x_midline(j))
        y_midline(j) = mean(y(x==x_midline(j)));
        width(j) = max(y(x==x_midline(j)))-min(y(x==x_midline(j)));
    end
    end

    if flag_fig
        figure(10)
        axis equal
        imshowpair(BW,BW1)
        hold all
        plot(x_midline, y_midline)
    end

    if isempty(Midline)
        stat.BW = (BW+BW1)/2;
        stat.nose = [x_nose,y_nose];
    end

    Midline = [Midline;y_midline];
    Width = [Width;width];

end
% center Midline to zero
Midline = Midline - (max(Midline,[],'all') + min(Midline,[],'all'))/2;

try
    stat.width = mean(Width,1,'omitnan')/2;
    stat.length = sum(~isnan(width));
    stat.midline = Midline/stat.length;
    stat.amplitude_nose = (max(Midline(:,1)) - min(Midline(:,1)))/stat.length;
    stat.amplitude_tail = (max(Midline,[],'all') - min(Midline,[],'all'))/stat.length;
catch 
    stat = [];
end
end

% This function generate a background by averaging N video frames
function Background = background_gen(vid,N,flag_debug)

    if nargin < 2 || isempty(N)
        N = 30;
    end
    
    if nargin < 3 
        flag_debug = false;
    end
        
    t = linspace(0,vid.Duration-1,N);
    
    vidCurrentTime = t(1);
    vidFrame = readFrame(vid);
    Background = double(vidFrame);
    for i = 2:N
        vidCurrentTime = t(i);
        vidFrame = readFrame(vid);
        Background = Background + double(vidFrame);
    end
    Background = Background./N/255;
    
    if flag_debug
        figure
        subplot(1,2,1)
        imshow(Background)
        title('background')
        subplot(1,2,2)
        imshow(vidFrame)
        title('frame')
        pause
    end
end



